{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Alphafold_single.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/Alphafold_single.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iVVz9evNDNOf"
      },
      "source": [
        "#AlphaFold_single"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "woIxeCPygt7K",
        "cellView": "form"
      },
      "source": [
        "#@title Install software\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "#@markdown on the left.\n",
        "\n",
        "# setup device\n",
        "import os\n",
        "import tensorflow as tf\n",
        "import jax\n",
        "\n",
        "try:\n",
        "  # check if TPU is available\n",
        "  import jax.tools.colab_tpu\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "  print('Running on TPU')\n",
        "  DEVICE = \"tpu\"\n",
        "except:\n",
        "  if jax.local_devices()[0].platform == 'cpu':\n",
        "    print(\"WARNING: no GPU detected, will be using CPU\")\n",
        "    DEVICE = \"cpu\"\n",
        "  else:\n",
        "    print('Running on GPU')\n",
        "    DEVICE = \"gpu\"\n",
        "    # disable GPU on tensorflow\n",
        "    tf.config.set_visible_devices([], 'GPU')\n",
        "\n",
        "    \n",
        "from IPython.utils import io\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar'\n",
        "PARAMS_DIR = './alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "# if not already installed\n",
        "try:\n",
        "  with tqdm.notebook.tqdm(total=50, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "      if not os.path.isdir(\"alphafold\"):\n",
        "        %shell rm -rf alphafold\n",
        "        %shell git clone {GIT_REPO} alphafold\n",
        "        %shell (cd alphafold; git checkout 1e216f93f06aa04aa699562f504db1d02c3b704c --quiet)\n",
        "\n",
        "        # colabfold patches\n",
        "        %shell mkdir --parents tmp\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/pairmsa.py\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/protein.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/config.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/model.patch -P tmp/\n",
        "        %shell wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/modules.patch -P tmp/\n",
        "\n",
        "        # Apply multi-chain patch from Lim Heo @huhlim\n",
        "        %shell patch -u alphafold/alphafold/common/protein.py -i tmp/protein.patch\n",
        "        \n",
        "        # Apply patch to dynamically control number of recycles (idea from Ryan Kibler)\n",
        "        %shell patch -u alphafold/alphafold/model/model.py -i tmp/model.patch\n",
        "        %shell patch -u alphafold/alphafold/model/modules.py -i tmp/modules.patch\n",
        "        %shell patch -u alphafold/alphafold/model/config.py -i tmp/config.patch\n",
        "        pbar.update(4)\n",
        "\n",
        "        %shell pip3 install ./alphafold\n",
        "        pbar.update(4)\n",
        "      \n",
        "        # speedup from kaczmarj\n",
        "        %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "        %shell curl -fsSL \"{SOURCE_URL}\" | tar x -C \"{PARAMS_DIR}\"\n",
        "        pbar.update(14+27)\n",
        "\n",
        "        # Install py3dmol.\n",
        "        %shell pip install py3dmol\n",
        "        pbar.update(1)\n",
        "      else:\n",
        "        pbar.update(50)\n",
        "\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "########################################################################################\n",
        "# --- Python imports ---\n",
        "import colabfold as cf\n",
        "import sys\n",
        "\n",
        "from google.colab import files\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "\n",
        "from alphafold.model import model\n",
        "from alphafold.model import config\n",
        "from alphafold.model import data\n",
        "from alphafold.data import pipeline\n",
        "from alphafold.common import protein\n",
        "\n",
        "from alphafold.model.tf import shape_placeholders\n",
        "\n",
        "def make_fixed_size(feat, shape_schema, num_res=100, msa_cluster_size=1, extra_msa_size=1, num_templates=0):\n",
        "  \"\"\"Guess at the MSA and sequence dimensions to make fixed size.\"\"\"\n",
        "  pad_size_map = {\n",
        "      shape_placeholders.NUM_RES: num_res,\n",
        "      shape_placeholders.NUM_MSA_SEQ: msa_cluster_size,\n",
        "      shape_placeholders.NUM_EXTRA_SEQ: extra_msa_size,\n",
        "      shape_placeholders.NUM_TEMPLATES: num_templates,\n",
        "  }\n",
        "  for k, v in feat.items():\n",
        "    # Don't transfer this to the accelerator.\n",
        "    if k == 'extra_cluster_assignment':\n",
        "      continue\n",
        "    shape = list(v.shape)\n",
        "    schema = shape_schema[k]\n",
        "    assert len(shape) == len(schema), (\n",
        "        f'Rank mismatch between shape and shape schema for {k}: '\n",
        "        f'{shape} vs {schema}')\n",
        "    pad_size = [pad_size_map.get(s2, None) or s1 for (s1, s2) in zip(shape, schema)]\n",
        "    padding = [(0, p - tf.shape(v)[i]) for i, p in enumerate(pad_size)]\n",
        "    if padding:\n",
        "      feat[k] = tf.pad(v, padding, name=f'pad_to_fixed_{k}')\n",
        "      feat[k].set_shape(pad_size)\n",
        "  return {k:np.asarray(v) for k,v in feat.items()}\n",
        "\n",
        "def parse_results(prediction_result, processed_feature_dict):\n",
        "  b_factors = prediction_result['plddt'][:,None] * prediction_result['structure_module']['final_atom_mask']  \n",
        "  dist_bins = jax.numpy.append(0,prediction_result[\"distogram\"][\"bin_edges\"])\n",
        "  dist_mtx = dist_bins[prediction_result[\"distogram\"][\"logits\"].argmax(-1)]\n",
        "  contact_mtx = jax.nn.softmax(prediction_result[\"distogram\"][\"logits\"])[:,:,dist_bins < 8].sum(-1)\n",
        "\n",
        "  out = {\"unrelaxed_protein\": protein.from_prediction(processed_feature_dict, prediction_result, b_factors=b_factors),\n",
        "         \"plddt\": prediction_result['plddt'],\n",
        "         \"pLDDT\": prediction_result['plddt'].mean(),\n",
        "         \"dists\": dist_mtx,\n",
        "         \"adj\": contact_mtx}\n",
        "\n",
        "  if \"ptm\" in prediction_result:\n",
        "    out.update({\"pae\": prediction_result['predicted_aligned_error'],\n",
        "                \"pTMscore\": prediction_result['ptm']})\n",
        "  return out\n",
        "\n",
        "if \"model_runner\" not in dir():\n",
        "  name = \"model_5_ptm\"\n",
        "  model_config = config.model_config(name)\n",
        "  model_config.data.eval.max_msa_clusters = 1\n",
        "  model_config.data.common.max_extra_msa = 1\n",
        "  model_config.model.num_recycle = 2\n",
        "  model_config.data.common.num_recycle = 2\n",
        "\n",
        "  model_runner = model.RunModel(model_config, data.get_model_haiku_params(name,'./alphafold/data'))\n",
        "\n",
        "  eval_cfg = model_runner.config.data.eval\n",
        "  crop_feats = {k:[None]+v for k,v in dict(eval_cfg.feat).items()}   \n",
        "\n",
        "  MAX_LEN = 100  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "source": [
        "#@title Enter the amino acid sequence to fold ⬇️\n",
        "import re\n",
        "\n",
        "# define sequence\n",
        "sequence = 'GGGGGGGGGGGGGGGGGGG' #@param {type:\"string\"}\n",
        "sequence = re.sub(\"[^A-Z]\", \"\", sequence.upper())\n",
        "\n",
        "msas, deletion_matrices = [],[]\n",
        "msas.append([sequence])\n",
        "deletion_matrices.append([[0]*len(sequence)])\n",
        "\n",
        "num_res = len(sequence)\n",
        "feature_dict = {**pipeline.make_sequence_features(sequence, 'test', num_res),\n",
        "                **pipeline.make_msa_features(msas, deletion_matrices=deletion_matrices)}\n",
        "\n",
        "if num_res > MAX_LEN:\n",
        "  MAX_LEN = num_res\n",
        "\n",
        "# process features\n",
        "processed_feature_dict = model_runner.process_features(feature_dict, random_seed=0)\n",
        "processed_feature_dict = make_fixed_size(processed_feature_dict, crop_feats, num_res=MAX_LEN)\n",
        "\n",
        "# predict\n",
        "prediction_result, (r, t) = model_runner.predict(processed_feature_dict)\n",
        "\n",
        "# save results\n",
        "outs = parse_results(prediction_result, processed_feature_dict)\n",
        "outs.update({\"recycles\":r, \"tol\":t})\n",
        "with open(\"out.pdb\", 'w') as f: f.write(protein.to_pdb(outs[\"unrelaxed_protein\"]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KAZj6CBZTkJM",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "color = \"chain\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "cf.show_pdb(\"out.pdb\", show_sidechains, show_mainchains, color,\n",
        "            color_HP=True, size=(800,480)).show()\n",
        "if color == \"lDDT\":\n",
        "  cf.plot_plddt_legend().show()  \n",
        "if \"pae\" in outs:\n",
        "  cf.plot_confidence(outs[\"plddt\"][:num_res], outs[\"pae\"][:num_res,:][:,:num_res]).show()\n",
        "else:\n",
        "  cf.plot_confidence(outs[\"plddt\"][:num_res]).show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}